import argparse
import logging

def get_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="Train RePaint DDPM for image inpainting")
    
    # 数据参数
    parser.add_argument("--data_dir", type=str, default="/root/code/datasets/Places365_256x256/", help="Dataset root directory")
    parser.add_argument("--image_size", type=int, default=256, help="Input image size")
    parser.add_argument("--batch_size", type=int, default=8, help="Training batch size")
    parser.add_argument("--eval_batch_size", type=int, default=8, help="Evaluation batch size")
    parser.add_argument("--num_workers", type=int, default=4, help="Number of data loader workers")
    parser.add_argument("--train_percentage", type=int, default=1, help="Percentage of training set to use")
    parser.add_argument("--val_percentage", type=int, default=1, help="Percentage of validation set to use")

    # 训练参数
    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate")
    parser.add_argument("--min_learning_rate", type=float, default=1e-6, help="Minimum learning rate")
    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay")
    parser.add_argument("--num_epochs", type=int, default=100, help="Number of training epochs")
    parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Gradient clipping threshold")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--device", type=str, default="cuda", help="Device to use")
    
    # 扩散模型参数
    parser.add_argument("--diffusion_steps", type=int, default=1000, help="Number of diffusion steps")
    parser.add_argument("--beta_start", type=float, default=1e-4, help="Starting beta value")
    parser.add_argument("--beta_end", type=float, default=0.02, help="Ending beta value")
    parser.add_argument("--time_embedding_dim", type=int, default=256, help="Time embedding dimension")
    
    # RePaint特定参数
    parser.add_argument("--num_resample_U", type=int, default=10, 
                       help="Number of resampling steps at each time step (U in paper)")
    parser.add_argument("--repaint_jump_length", type=int, default=10, 
                       help="Jump length for resampling schedule (j in paper)")
    parser.add_argument("--repaint_jump_n_sample", type=int, default=10, 
                       help="Number of resampling at each jump point (r in paper)")
    parser.add_argument("--repaint_jump_start_fraction", type=float, default=0.0,
                       help="Fraction of T to start jump resampling (0.0 means start from beginning)")
    parser.add_argument("--use_repaint_schedule_approx", action="store_true", default=True,
                       help="Use approximate schedule from Figure 9 of RePaint paper")
    
    # 保存和可视化参数
    parser.add_argument("--output_dir", type=str, default="/root/code/My/InpaintingDDPM/outputs", help="Output directory")
    parser.add_argument("--vis_interval", type=int, default=5000, help="Steps interval for visualizing results")
    parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="Resume training from checkpoint")
    
    # 评估参数
    parser.add_argument("--do_eval", action="store_false", help="Whether to perform evaluation")
    parser.add_argument("--eval_interval", type=int, default=1, help="Epochs interval for evaluation")

    # 早停
    parser.add_argument('--early_stopping_patience', type=int, default=10, 
                   help='连续多少个评估周期验证损失不下降就触发早停')
    
    return parser.parse_args()